import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import glob
import os, os.path
import math
import pickle
import scipy.sparse as sp
import random
from random import choice

#二维网格模型
#row_length 行长  col_length 列长
def grid_2d(row_length,col_length):
    length = row_length+col_length
    G = nx.grid_2d_graph(row_length, col_length, periodic=False)
    oriadj = nx.adjacency_matrix(G)
    oriadj_= oriadj.toarray()
    nodes_list = list(G.nodes())
    ori_nodesNum = len(G.edges)
    #addEdgeNum随机加边数量
    # addEdgeNum = int(row_length*random.random()+col_length*random.random())
    count_add = 0
    count_remove = 0
    orig_num_cc = nx.number_connected_components(G)
    news_edges = []
    print(len(G.edges))
    while(count_add<int(ori_nodesNum/5)):
          node1 = choice(nodes_list)
          node2 = choice(nodes_list)
          if(G.has_edge(node1,node2)):
              continue
          #距离越短，加边概率越高
          # if(random.random() < (1.0-((abs(node2[0] - node1[0]) + abs(node2[1] - node1[1]))/length)) and node1!=node2):
          if(random.random() < (1.0-(nx.shortest_path_length(G,node1,node2)/length)) and node1!=node2):
              print("\tnew edge:\t {} -- {} -- ".format(node1, node2))
              print("dis:{}".format(nx.shortest_path_length(G,node1,node2)))
              G.add_edge(node1, node2)
              news_edges.append((node1[0]*col_length+node1[1],node2[0]*col_length+node2[1]))
              count_add = count_add+1
    adj = nx.adjacency_matrix(G)
    adj_ = adj.toarray()
    # for node in G.nodes():
    #       new = choice(list(G.neighbors(node)))
    #       if(random.random()>0.5 and G.has_edge(new,node)):
    #           G.remove_edge(new,node)
    #           if(nx.number_connected_components(G) > orig_num_cc):
    #               G.add_edge(new,node)
    #               continue
    #           count_remove = count_remove+1
    print("add-{}-remove-{}".format(count_add,count_remove))
    print(len(G.edges))
    for pos in news_edges:
        if(adj_[pos[0]][pos[1]]!=1 or adj_[pos[1]][pos[0]]!=1):
            print("error")
        adj_[pos[0]][pos[1]] = 0
        adj_[pos[1]][pos[0]] = 0
    for i in range(row_length):
        for j in range(col_length):
            if(adj_[i][j]!=oriadj_[i][j]):
                print("error")
    # nx.draw(G)
    # plt.show()
    return G,news_edges
if __name__ == '__main__':
    grid_2d(20,20)

